import json
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
import torch
from torch.utils.data import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer
)


# 1. 数据加载与预处理
def load_data(file_path):
    texts = []
    labels = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            data = json.loads(line)
            texts.append(data['text'])
            labels.append(data['label'])
    return texts, labels


def load_test_data(file_path):
    texts = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            data = json.loads(line)
            texts.append(data['text'])
    return texts


# 2. 数据集类定义
class TextDataset(Dataset):
    def __init__(self, encodings, labels=None):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        if self.labels is not None:
            item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.encodings['input_ids'])


# 3. 主函数
def main():
    # 加载数据
    train_texts, train_labels = load_data('train.jsonl')
    test_texts = load_test_data('test.jsonl')

    # 划分训练集和验证集
    train_texts, val_texts, train_labels, val_labels = train_test_split(
        train_texts, train_labels, test_size=0.1, random_state=42
    )

    # 初始化tokenizer
    model_name = "bert-base-uncased"
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # 数据预处理
    def preprocess_function(examples):
        return tokenizer(examples, truncation=True, padding='max_length', max_length=256)

    train_encodings = preprocess_function(train_texts)
    val_encodings = preprocess_function(val_texts)
    test_encodings = preprocess_function(test_texts)

    # 创建数据集
    train_dataset = TextDataset(train_encodings, train_labels)
    val_dataset = TextDataset(val_encodings, val_labels)
    test_dataset = TextDataset(test_encodings)

    # 加载模型
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=2,
        ignore_mismatched_sizes=True  # 忽略尺寸不匹配的警告
    )

    # 定义评估指标
    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        predictions = np.argmax(logits, axis=-1)
        return {'f1': f1_score(labels, predictions, average='binary')}

    # 训练参数 - 更新为兼容旧版本的参数名
    training_args = TrainingArguments(
        output_dir='./results',
        eval_strategy="epoch",  # 旧版本使用 eval_strategy 而不是 evaluation_strategy
        save_strategy="epoch",
        learning_rate=2e-5,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=16,
        num_train_epochs=3,
        weight_decay=0.01,
        load_best_model_at_end=True,
        metric_for_best_model='f1',
        logging_dir='./logs',
        logging_steps=100,
        disable_tqdm=False  # 显示进度条
    )

    # 创建Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        compute_metrics=compute_metrics,
    )

    # 训练模型
    trainer.train()

    # 评估模型
    eval_results = trainer.evaluate()
    print(f"Validation F1 score: {eval_results['eval_f1']:.4f}")

    # 预测测试集
    predictions = trainer.predict(test_dataset)
    pred_labels = np.argmax(predictions.predictions, axis=-1)

    # 保存预测结果
    with open('submit.txt', 'w') as f:
        for label in pred_labels:
            f.write(f"{label}\n")


if __name__ == "__main__":
    main()